import os
import sys
import numpy as np
import torch
import platform
import os

def getGPU():
    # 获取每个 GPU 的剩余显存数，并存放到 tmp 文件中
    if platform.system().lower() == 'windows':
        return 0
    elif platform.system().lower() == 'linux':
        if torch.cuda.device_count() >1:
            os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp')  #读取gpu memory
            memory_gpu = [int(x.split()[2]) for x in open('tmp', 'r').readlines()]
            print("剩余显存最大的 GPU 编号是： ",np.argmax(memory_gpu))  # 获取剩余显存最多的 GPU 的编号
            os.system('rm tmp')  # 删除临时生成的 tmp 文件
            return np.argmax(memory_gpu)  #返回GPU中剩余显存最多的序号
        else:
            return 0



if __name__=='__main__':
    print(sys.path)

    getGPU()
    device=None
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    else:
        torch.device(device)
    print(device)